package rmmk.framework.datasources;

import java.util.ArrayList;
import java.util.List;
import java.util.Vector;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import rmmk.datasources.parsing.ClassTypes;

/**
 * Klasa umożliwiająca selekcję dokumentów ze zbioru na podstawie kryteriów
 * takich jak: ilość i rodzaj posiadanych kategorii. Dzieli również zbiór
 * dokumentów na testowy i treningowy
 * 
 * @author marcin
 * 
 */
public class DocumentSetSelector {

	private static Logger logger = LoggerFactory
			.getLogger(DocumentSetSelector.class);
	private static List<ClassTypes> desiredCategories = new ArrayList<>();
	private int limit = 0;
	private int lowerLimit = 0;
	private List<Document> allDocuments;
	private double trainToTestSetProportion = 0.6;
	private ArrayList<String> avilableCategoryNames = new ArrayList<>();

	public DocumentSetSelector(List<Document> allDocuments) {
		this.allDocuments = allDocuments;
	}

	private List<Document> selectDocuments() {
		Vector<Document> selectedDocuments = new Vector<>();

		if (lowerLimit == 0) {
			for (Document d : this.allDocuments) {
				boolean containsAll = true;

				for (int i = 0; i < desiredCategories.size(); ++i) {
					if (!d.containsCategory(desiredCategories.get(i), limit,
							this.avilableCategoryNames)) {
						containsAll = false;
						break;
					}
				}

				if (containsAll) {
					selectedDocuments.add(d);
				}
			}
		} else {
			for (Document d : this.allDocuments) {
				boolean containsAny = false;

				for (int i = 0; i < desiredCategories.size(); ++i) {
					if (d.containsCategoryLowerLimit(desiredCategories.get(i), lowerLimit)) {
						containsAny = true;
						break;
					}
				}

				if (containsAny) {
					selectedDocuments.add(d);
				}
			}
		}

		logger.info("Wyselekcjonowano dokumentów: " + selectedDocuments.size()
				+ " z " + allDocuments.size());
		return selectedDocuments;
	}

	public List<Document> getTrainSet() {
		List<Document> selected = this.selectDocuments();
		List<Document> ret = new Vector<>();

		for (int i = 0; i < selected.size() * this.trainToTestSetProportion; ++i)
			ret.add(selected.get(i));

		logger.info("Dokumentów treningowych: " + ret.size());
		return ret;
	}

	public List<Document> getTestSet() {
		List<Document> selected = this.selectDocuments();
		List<Document> trainDocuments = this.getTrainSet();

		for (int i = 0; i < trainDocuments.size(); ++i)
			selected.remove(trainDocuments.get(i));

		logger.info("Dokumentów testowych: " + selected.size());
		return selected;
	}

	public static List<ClassTypes> getDesiredCategories() {
		return desiredCategories;
	}

	public void setDesiredCategories(List<ClassTypes> desiredCategories) {
		DocumentSetSelector.desiredCategories = desiredCategories;
	}

	public void addDesiredCategory(ClassTypes desiredCategory) {
		DocumentSetSelector.desiredCategories.add(desiredCategory);
	}

	public int getLimit() {
		return limit;
	}

	public void setLimit(int limit) {
		this.limit = limit;
	}

	public List<Document> getAllDocuments() {
		return allDocuments;
	}

	public void setAllDocuments(List<Document> allDocuments) {
		this.allDocuments = allDocuments;
	}

	public double getTrainToTestSetProportion() {
		return trainToTestSetProportion;
	}

	public void setTrainToTestSetProportion(double trainToTestSetProportion) {
		this.trainToTestSetProportion = trainToTestSetProportion;
	}

	public ArrayList<String> getAvilableCategoryNames() {
		return avilableCategoryNames;
	}

	public void setAvilableCategoryNames(ArrayList<String> avilableCategoryNames) {
		this.avilableCategoryNames = avilableCategoryNames;
	}

	public void addCategoryName(String name) {
		this.avilableCategoryNames.add(name);
	}

	public int getLowerLimit() {
		return lowerLimit;
	}

	public void setLowerLimit(int lowerLimit) {
		this.lowerLimit = lowerLimit;
	}
}
